# 全局变量
global_trace_enabled = False  # 用于控制是否启用钩子逻辑
trace_logs = []  # 记录钩子函数执行的日志

# 示例目标函数
class math_ops:
    @staticmethod
    def add(a, b):
        return a + b

def register_node(fn_path):
    """
    注册一个钩子函数到指定函数路径上
    """
    # 假设函数在全局命名空间中，用 eval 获取
    fn_names = fn_path.split(".")
    fn_module = eval(".".join(fn_names[:-1]))
    fn_name = fn_names[-1]

    # 保存原始函数
    original_fn = getattr(fn_module, fn_name)

    def make_hook(bind_fn):
        """
        生成一个钩子函数，用于在原始函数前后注入逻辑
        """
        call_count = 0  # 用于标记函数调用的次数

        def hook_function(*args, **kwargs):
            """
            内部钩子函数，取代原始函数执行
            """
            nonlocal call_count
            global global_trace_enabled

            # 如果未启用追踪，直接调用原始函数
            if not global_trace_enabled:
                return original_fn(*args, **kwargs)

            # 启用追踪时，记录调用过程
            trace_logs.append(f"Calling {fn_path} with args={args}, kwargs={kwargs}")

            # 调用原始函数
            result = original_fn(*args, **kwargs)

            # 执行绑定的逻辑
            bind_fn(call_count, result, *args, **kwargs)

            # 记录完成调用后的结果
            trace_logs.append(f"Result from {fn_path}: {result}")

            # 增加调用计数
            call_count += 1
            return result

        # 替换目标函数为钩子版本
        setattr(fn_module, fn_name, hook_function)

    return make_hook


# 简单的绑定逻辑
@register_node("math_ops.add")
def add_hook(call_count, result, a, b):
    """
    钩子函数逻辑：记录每次 `math_ops.add` 被调用的信息
    """
    trace_logs.append(f"[HOOK] add() called #{call_count}: {a} + {b} = {result}")





# 测试逻辑
if __name__ == "__main__":
    global_trace_enabled = True  # 启用钩子逻辑

    # 调用目标函数
    print(math_ops.add(3, 5))
    print(math_ops.add(10, 20))

    # 查看日志
    global_trace_enabled = False  # 关闭钩子逻辑
    print("\nTrace logs:")
    for log in trace_logs:
        print(log)
